import os
import argparse
from PIL import Image
import piexif
from datetime import datetime

def set_capture_time(image_path, new_time):
    try:
        exif_dict = piexif.load(image_path)
        if piexif.ImageIFD.DateTime in exif_dict['0th']:
            print(f"Image {image_path} already has a capture time: {exif_dict['0th'][piexif.ImageIFD.DateTime].decode()}")
            return False
        
        # Convert the string time to the format required for EXIF
        time_to_set = datetime.strptime(new_time, '%Y-%m-%d %H:%M:%S').strftime('%Y:%m:%d %H:%M:%S')
        
        exif_dict['0th'][piexif.ImageIFD.DateTime] = time_to_set
        exif_dict['Exif'][piexif.ExifIFD.DateTimeOriginal] = time_to_set
        exif_dict['Exif'][piexif.ExifIFD.DateTimeDigitized] = time_to_set
        
        exif_bytes = piexif.dump(exif_dict)
        img = Image.open(image_path)
        img.save(image_path, 'jpeg', exif=exif_bytes)
        print(f"Set capture time of {image_path} to {new_time}")
        return True
    except Exception as e:
        print(f"Failed to set capture time for {image_path}: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description="Fix or set the capture time of photos in a directory.")
    parser.add_argument('--scan-dir', required=True, help='Directory containing photos.')
    parser.add_argument('--time', required=True, help='The preset capture time in the format "YYYY-MM-DD HH:MM:SS".')
    
    args = parser.parse_args()
    
    for root, _, files in os.walk(args.scan_dir):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                image_path = os.path.join(root, file)
                set_capture_time(image_path, args.time)

if __name__ == "__main__":
    main()